import sys
import time
import torch

from ai import utils
from ai.run.base_train_runner import GxlBaseRunner


class RunnerGxl(GxlBaseRunner):
    def __init__(self, model, optim, loss_f, train_loader,
                 logger, valid_loader=None, scheduler=None, multi=False,
                 local_rank=0, is_class=True,
                 device=torch.device('cpu')):
        super(RunnerGxl, self).__init__(model, optim, loss_f, train_loader,
                                        logger, valid_loader, scheduler, multi,
                                        local_rank, is_class, device)

    def train_function(self, epochs):
        for epoch in range(epochs):
            train_loss = 0.0
            batch_size = 0
            st = time.time()
            for X, y in self.train_loader:
                self.model.train()
                self.optim.zero_grad()
                X, y = utils.utils_model.put_data_to_device(X, y, self.device)
                out = self.model(X)
                loss = self.loss_f(out, y)
                loss.backward()
                self.optim.step()
                if self.scheduler is not None:
                    self.scheduler.step()
                train_loss += loss.item()
                batch_size += 1
            train_loss = train_loss / batch_size
            valid_loss = self.calculate_valid_loss()
            et = time.time()
            if self.local_rank == 0:
                self.logger.info(
                    f'epoch{epoch}:train_loss:{train_loss:.4f};valid_loss:{valid_loss:.4f};;time:{et - st:.2f}s')


if __name__ == '__main__':
    """"""
